from libs.PipeLine import PipeLine, ScopedTiming
from libs.AIBase import AIBase
from libs.AI2D import Ai2d
import os
import ujson
from media.media import *
import nncase_runtime as nn
import ulab.numpy as np
import image
import gc
import sys
import aidemo
from machine import Pin
from machine import FPIOA
import time
import math

class PersonKeyPointApp(AIBase):
    # 保持原有代码不变
    def __init__(self, kmodel_path, model_input_size, confidence_threshold=0.1, nms_threshold=0.5, rgb888p_size=[1280, 720], display_size=[1920, 1080], debug_mode=0):
        super().__init__(kmodel_path, model_input_size, rgb888p_size, debug_mode)
        self.kmodel_path = kmodel_path
        self.model_input_size = model_input_size
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.rgb888p_size = [ALIGN_UP(rgb888p_size[0], 16), rgb888p_size[1]]
        self.display_size = [ALIGN_UP(display_size[0], 16), display_size[1]]
        self.debug_mode = debug_mode
        self.SKELETON = [(16, 14), (14, 12), (17, 15), (15, 13), (12, 13), (6, 12), (7, 13), (6, 7), (6, 8), (7, 9), (8, 10), (9, 11), (2, 3), (1, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 7)]
        self.LIMB_COLORS = [(255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 255, 51, 255), (255, 255, 51, 255), (255, 255, 51, 255), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0)]
        self.KPS_COLORS = [(255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255)]
        self.ai2d = Ai2d(debug_mode)
        self.ai2d.set_ai2d_dtype(nn.ai2d_format.NCHW_FMT, nn.ai2d_format.NCHW_FMT, np.uint8, np.uint8)

    def config_preprocess(self, input_image_size=None):
        with ScopedTiming("set preprocess config", self.debug_mode > 0):
            ai2d_input_size = input_image_size if input_image_size else self.rgb888p_size
            top, bottom, left, right = self.get_padding_param()
            self.ai2d.pad([0, 0, 0, 0, top, bottom, left, right], 0, [0, 0, 0])
            self.ai2d.resize(nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
            self.ai2d.build([1, 3, ai2d_input_size[1], ai2d_input_size[0]], [1, 3, self.model_input_size[1], self.model_input_size[0]])

    def postprocess(self, results):
        with ScopedTiming("postprocess", self.debug_mode > 0):
            results = aidemo.person_kp_postprocess(results[0], [self.rgb888p_size[1], self.rgb888p_size[0]], self.model_input_size, self.confidence_threshold, self.nms_threshold)
            return results

    def draw_result(self, pl, res):
        with ScopedTiming("display_draw", self.debug_mode > 0):
            if res[0]:
                pl.osd_img.clear()
                kpses = res[1]
                for i in range(len(res[0])):
                    for k in range(17 + 2):
                        if k < 17:
                            kps_x, kps_y, kps_s = round(kpses[i][k][0]), round(kpses[i][k][1]), kpses[i][k][2]
                            kps_x1 = int(float(kps_x) * self.display_size[0] // self.rgb888p_size[0])
                            kps_y1 = int(float(kps_y) * self.display_size[1] // self.rgb888p_size[1])
                            if kps_s > 0:
                                pl.osd_img.draw_circle(kps_x1, kps_y1, 5, self.KPS_COLORS[k], 4)
                        ske = self.SKELETON[k]
                        pos1_x, pos1_y = round(kpses[i][ske[0]-1][0]), round(kpses[i][ske[0]-1][1])
                        pos1_x_ = int(float(pos1_x) * self.display_size[0] // self.rgb888p_size[0])
                        pos1_y_ = int(float(pos1_y) * self.display_size[1] // self.rgb888p_size[1])

                        pos2_x, pos2_y = round(kpses[i][(ske[1]-1)][0]), round(kpses[i][(ske[1]-1)][1])
                        pos2_x_ = int(float(pos2_x) * self.display_size[0] // self.rgb888p_size[0])
                        pos2_y_ = int(float(pos2_y) * self.display_size[1] // self.rgb888p_size[1])

                        pos1_s, pos2_s = kpses[i][(ske[0]-1)][2], kpses[i][(ske[1]-1)][2]
                        if pos1_s > 0.0 and pos2_s > 0.0:
                            pl.osd_img.draw_line(pos1_x_, pos1_y_, pos2_x_, pos2_y_, self.LIMB_COLORS[k], 4)
                    pl.osd_img.draw_circle(kps_x1, kps_y1, 3, (0, 255, 0), 2)  # 绿色圆心表示肩臀中点
                    gc.collect()
            else:
                pl.osd_img.clear()

    def get_padding_param(self):
        dst_w = self.model_input_size[0]
        dst_h = self.model_input_size[1]
        input_width = self.rgb888p_size[0]
        input_high = self.rgb888p_size[1]
        ratio_w = dst_w / input_width
        ratio_h = dst_h / input_high
        ratio = min(ratio_w, ratio_h)
        new_w = int(ratio * input_width)
        new_h = int(ratio * input_high)
        dw = (dst_w - new_w) // 2
        dh = (dst_h - new_h) // 2
        return dh, dh, dw, dw

class MultiBehaviorDetector:
    def __init__(self, display_size, rgb888p_size, debug_mode=3):  # 新增debug_mode参数
        self.display_size = display_size
        self.rgb888p_size = rgb888p_size
        self.min_confidence = 0.4
        
        # 新增参数
        self.crowd_distance_threshold = 1.2  # 聚集距离阈值
        self.crowd_centers_threshold = 0.3  # 聚集对数阈值
        
        # 调试模式
        self.debug_mode = debug_mode
        # 聚集检测冷却机制
        self.last_crowd_detected_time = 0  # 上次检测到聚集的时间戳
        self.crowd_cooldown_seconds = 4    # 冷却时间（秒）

    
    # 新增：通用关键点置信度检查
    def _check_keypoint_confidence(self, keypoints, indices):
        """检查指定索引的关键点置信度是否都达标"""
        for idx in indices:
            if keypoints[idx][2] < self.min_confidence:
                # if self.debug_mode > 1:
                    # print(f"关键点 {idx} 置信度不足: {keypoints[idx][2]} < {self.min_confidence}")
                return False
        return True
    
    # def get_view_type(self, keypoints):
    #     """判断视图类型（正视图或侧视图）"""
    #     if not self._check_keypoint_confidence(keypoints, [5, 6]):  # 左右肩
    #         return None
            
    #     left_shoulder = keypoints[5]
    #     right_shoulder = keypoints[6]
    #     shoulder_distance = math.sqrt((right_shoulder[0] - left_shoulder[0])**2 + 
    #                                  (right_shoulder[1] - left_shoulder[1])**2)
    #     return "front" if shoulder_distance > 50 else "side"

    # def get_reference_distance(self, keypoints):
    #     """根据视图类型获取参考距离"""
    #     view_type = self.get_view_type(keypoints)
    #     if view_type is None:
    #         return 0
            
    #     left_shoulder = keypoints[5]
    #     right_shoulder = keypoints[6]
    #     shoulder_distance = math.sqrt((right_shoulder[0] - left_shoulder[0])**2 + 
    #                                  (right_shoulder[1] - left_shoulder[1])**2)
                                     
    #     return shoulder_distance * 2 if view_type == "front" else shoulder_distance * 1.5




    def detect_crowd(self, all_keypoints, current_time):
        """改进版聚集检测算法，添加冷却机制"""
        # 检查是否在冷却期内
        if current_time - self.last_crowd_detected_time < self.crowd_cooldown_seconds:
            # 处于冷却期，直接返回False，不进行聚集检测
            if self.debug_mode > 1:
                print(f"处于聚集检测冷却期，剩余 {self.crowd_cooldown_seconds - (current_time - self.last_crowd_detected_time):.1f} 秒")
            return False
            
        centers = []
        
        if len(all_keypoints) > 3:
            # 计算每个人的中心点和身体尺寸
            for kp in all_keypoints:
                # 检查关键点置信度
                if not self._check_keypoint_confidence(kp, [5, 6, 11, 12]):
                    continue
                    
                left_shoulder, right_shoulder = kp[5], kp[6]
                left_hip, right_hip = kp[11], kp[12]
                
                # 计算身体中心
                center_x = (left_shoulder[0] + right_shoulder[0] + left_hip[0] + right_hip[0]) / 4
                center_y = (left_shoulder[1] + right_shoulder[1] + left_hip[1] + right_hip[1]) / 4
                
                # 计算身体宽度和高度
                width = max(abs(right_shoulder[0] - left_shoulder[0]), 
                        abs(right_hip[0] - left_hip[0]))
                height = max(abs(right_shoulder[1] - right_hip[1]), 
                            abs(left_shoulder[1] - left_hip[1]))
                reference_size = (width + height) / 2
                
                centers.append((center_x, center_y, reference_size))
            
            # 计算人群密度
            crowd_count = 0
            for i in range(len(centers)):
                for j in range(i + 1, len(centers)):
                    dx = centers[i][0] - centers[j][0]
                    dy = centers[i][1] - centers[j][1]
                    distance = math.sqrt(dx**2 + dy**2)
                    avg_size = (centers[i][2] + centers[j][2]) / 2
                    
                    if distance < avg_size * self.crowd_distance_threshold:
                        crowd_count += 1
            
            # 判断是否聚集
            threshold = len(centers) * self.crowd_centers_threshold
            is_crowd = crowd_count > threshold
            
            # 如果检测到聚集，更新最后检测时间
            if is_crowd:
                self.last_crowd_detected_time = current_time
                if self.debug_mode > 0:
                    print(f"检测到人群聚集，将冷却 {self.crowd_cooldown_seconds} 秒")
            
            if len(centers)>3:
                return is_crowd
            else:
                False
        else:
            return False


    def detect_fall(self, keypoints, current_time):
        """改进版跌倒检测算法"""
        # 提取关键点并检查置信度
        if not self._check_keypoint_confidence(keypoints, [6, 12, 15]):  # 右肩、右髋、右脚
            return False
            
        right_shoulder, right_hip, right_foot = keypoints[6], keypoints[12], keypoints[15]
        
        # 计算身体倾斜角度
        dx_shoulder_hip = right_hip[0] - right_shoulder[0]
        dy_shoulder_hip = right_hip[1] - right_shoulder[1]
        
        dx_hip_foot = right_foot[0] - right_hip[0]
        dy_hip_foot = right_foot[1] - right_hip[1]
        
        # 避免除零错误
        shoulder_hip_slope = float('inf') if dx_shoulder_hip == 0 else dy_shoulder_hip / dx_shoulder_hip
        hip_foot_slope = float('inf') if dx_hip_foot == 0 else dy_hip_foot / dx_hip_foot
            
        # 检查身体是否接近水平
        is_horizontal = (abs(shoulder_hip_slope) < 0.5) or (abs(hip_foot_slope) < 0.5)
        
        # 检查双脚是否接近肩部高度
        foot_shoulder_distance = abs(right_foot[1] - right_shoulder[1])
        shoulder_width = abs(keypoints[5][0] - keypoints[6][0])
        
        is_low = foot_shoulder_distance < shoulder_width * 0.7
        
        is_fall = is_horizontal and is_low
        
        if self.debug_mode > 0 and is_fall:
            print(f"检测到跌倒: 肩髋斜率={shoulder_hip_slope:.2f}, 髋脚斜率={hip_foot_slope:.2f}, 脚肩距={foot_shoulder_distance}")
            
        return is_fall
    
        
    def _check_keypoint_confidence(self, keypoints, indices):
        """检查指定关键点的置信度是否足够"""
        for idx in indices:
            if len(keypoints[idx]) < 3 or keypoints[idx][2] < 0.3:  # 置信度阈值
                return False
        return True

def init_ai_model(kmodel_path, rgb888p_size, display_size):
    person_kp = PersonKeyPointApp(
        kmodel_path,
        model_input_size=[320, 320],
        confidence_threshold=0.1,
        nms_threshold=0.5,
        rgb888p_size=rgb888p_size,
        display_size=display_size
    )
    person_kp.config_preprocess()
    
    detector = MultiBehaviorDetector(display_size, rgb888p_size)
    return person_kp, detector

# 执行AI推理
def run_inference(person_kp, detector, pl, current_time):
    img = pl.get_frame()
    res = person_kp.run(img)

    detection_results = {
        "fall": False,
        "run": False,
        "crowd": False,
        "fight": False  # 新增打架检测结果
    }

    if res[0] and len(res[1]) > 0:
        all_keypoints = res[1]

        # 单人检测
        for keypoints in all_keypoints:
            if detector.detect_fall(keypoints, current_time):
                detection_results["fall"] = True
                print("⚠️ 检测到跌倒")

        # 多人检测
        if len(all_keypoints) >= 2:
            if detector.detect_crowd(all_keypoints, current_time):
                detection_results["crowd"] = True
                print("⚠️ 检测到聚集")
    print(f"剩余内存: {gc.mem_free()}")
    person_kp.draw_result(pl, res)
    return detection_results